#include "StdAfx.h"
#include "user_token.h"
#include <WtsApi32.h>

#pragma comment(lib, "Wtsapi32.lib")

using namespace ATL;

//-------------------------------------------------------------------------------
user_token::user_token( HANDLE htoken )
{
	_token.Attach(htoken);
}
//-------------------------------------------------------------------------------
HANDLE user_token::get_token(void) const
{
	return _token.GetHandle();
}
//-------------------------------------------------------------------------------
security_identifier user_token::get_owner(void) const
{
	CSid sid;
	_token.GetOwner(&sid);
	return security_identifier(sid);
}
//-------------------------------------------------------------------------------
void user_token::set_owner(const security_identifier& si)
{
	if(!_token.SetOwner((const CSid&)si)){
		throw ::GetLastError();
	}
}
//-------------------------------------------------------------------------------
security_identifier user_token::get_logon(void) const
{
	CSid sid;
	_token.GetLogonSid(&sid);
	return security_identifier(sid);
}
//-------------------------------------------------------------------------------
security_identifier user_token::get_primary_group(void) const
{
	CSid sid;
	_token.GetPrimaryGroup(&sid);
	return security_identifier(sid);
}
//-------------------------------------------------------------------------------
void user_token::set_primary_group(const security_identifier& si)
{
	if(!_token.SetPrimaryGroup((const CSid&) si)){
		throw ::GetLastError();
	}
}
//-------------------------------------------------------------------------------
std::vector<security_identifier> user_token::get_groups(void) const
{
	std::vector<security_identifier> res;

	CTokenGroups groups;
	_token.GetGroups(&groups);

	CSid::CSidArray sids;
	groups.GetSidsAndAttributes(&sids, NULL);

	int len = sids.GetCount();
	for(int i=0 ; i<len ; i++){
		res.push_back(sids[i]);
	}

	return res;
}
//-------------------------------------------------------------------------------
bool user_token::get_is_primary_token(void) const
{
	TOKEN_TYPE tt;
	_token.GetType(&tt);

	return tt == TokenPrimary;
}
//-------------------------------------------------------------------------------
bool user_token::get_is_impersonation_token(void) const
{
	TOKEN_TYPE tt;
	_token.GetType(&tt);

	return tt == TokenImpersonation;
}
//-------------------------------------------------------------------------------
user_token::operator HANDLE()
{
	return _token.GetHandle();
}
//-------------------------------------------------------------------------------
user_token::operator ATL::CAccessToken&()
{
	return _token;
}
//-------------------------------------------------------------------------------
user_token user_token::get_from_session_id( UINT session_id )
{
	HANDLE h = NULL;
	if(!::WTSQueryUserToken(session_id, &h)){
		throw ::GetLastError();
	}

	return user_token(h);
}
//-------------------------------------------------------------------------------
user_token user_token::get_from_process( DWORD process_id, DWORD access /*= TOKEN_ALL_ACCESS*/ )
{
	HANDLE restoken = NULL;

	HANDLE hproc = ::OpenProcess(PROCESS_QUERY_INFORMATION, FALSE, process_id);
	if(!hproc){
		throw ::GetLastError();
	}

	if(!::OpenProcessToken(hproc, access, &restoken))
	{
		DWORD err = ::GetLastError();
		::CloseHandle(hproc);
		
		throw err;
	}

	::CloseHandle(hproc);
	return user_token(restoken);
}
//-------------------------------------------------------------------------------
user_token user_token::get_from_thread( DWORD thread_id, DWORD access /*= TOKEN_ALL_ACCESS*/ )
{
	HANDLE restoken = NULL;

	HANDLE hthread = ::OpenThread(THREAD_QUERY_INFORMATION, FALSE, thread_id);
	if(!hthread){
		throw ::GetLastError();
	}

	if(!::OpenThreadToken(hthread, access, FALSE, &restoken))
	{
		DWORD err = ::GetLastError();
		::CloseHandle(hthread);

		throw err;
	}

	::CloseHandle(hthread);
	return user_token(restoken);
}
//-------------------------------------------------------------------------------
HANDLE user_token::detach( void )
{
	return _token.Detach();
}
//-------------------------------------------------------------------------------
user_token get_from_login( const std::wstring& username, const std::wstring& domain /*= L""*/, const std::wstring& password /*= L""*/, DWORD logon_type /*= LOGON32_LOGON_INTERACTIVE*/, DWORD logon_provider /*= LOGON32_PROVIDER_DEFAULT*/)
{
	CAccessToken token;

	token.LogonUserW(username.c_str(),
						domain.empty() ? NULL : domain.c_str(),
						password.empty() ? NULL : password.c_str(),
						logon_type,
						logon_provider);

	return user_token(token.Detach());
}
//-------------------------------------------------------------------------------
